from ultralytics import YOLO
import time
import cv2
import os
import numpy as np

model = YOLO(
    '/home/yimu/Backup/projects_zy/yolov8/runs/detect/train63/weights/last.pt')  # load a pretrained model (recommended for training)
model.info()


def image_resize(image, resize_rate):
    src_w = image.shape[0]
    src_h = image.shape[1]

    # image = cv.resize(image, None, fx=resize_rate, fy=resize_rate, interpolation=cv.INTER_LINEAR)
    new_image = cv2.resize(image, None, fx=resize_rate, fy=resize_rate, interpolation=cv2.INTER_NEAREST)
    return new_image


def Watershed(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    # Apply the thresholding to create a binary image
    # print(cv2.THRESH_BINARY, cv2.THRESH_BINARY_INV, cv2.THRESH_OTSU, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU, cv2.THRESH_TRIANGLE)
    # ret, thres = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
    ret, thres = cv2.threshold(gray, 0, 255, cv2.THRESH_OTSU)
    # Perform a distance transform
    distance = cv2.distanceTransform(thres, cv2.DIST_L2, 5)
    ret, sure_fg = cv2.threshold(distance, 0.7 * distance.max(), 255, 0)
    # Perform the watershed algorithm
    sure_fg = np.uint8(sure_fg)
    unknown = cv2.subtract(thres, sure_fg)
    ret, markers = cv2.connectedComponents(sure_fg)
    # Add one to all labels so that sure background is not 0, but 1
    markers = markers + 1
    # Now, mark the region of unknown with zero　　
    markers[unknown == 255] = 0
    markers = cv2.watershed(img, markers)
    # Create the output image　　
    img[markers == -1] = [255, 0, 0]
    return gray


# def canny(img, min_thres, max_thres):

def list_dir(path, file_list, file_type='txt', filter_str=None):
    for file in os.listdir(path):
        file_path = os.path.join(path, file)
        if os.path.isdir(file_path):
            list_dir(file_path, file_list, file_type, filter_str=filter_str)
        else:
            if filter_str is None:
                if file_path[-len(file_type):].lower() == file_type.lower():
                    file_list.append(file_path)
            else:
                if filter_str not in file_path:
                    if file_path[-len(file_type):].lower() == file_type.lower():
                        file_list.append(file_path)


if __name__ == '__main__':

    file_path = '/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test'
    file_list = []
    list_dir(file_path, file_list, 'jpg')
    count = 0
    for img_path in file_list:
        # result = model.predict(source='/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test',save=False)
        #
        result = model(img_path)
        img = cv2.imread(img_path)
        for r in result:

            # print(r.boxes)
            for re in r.boxes:
                xyxy = re.cpu().numpy().xyxy.astype(int)
                crop_img = img[xyxy[0][1]:xyxy[0][3], xyxy[0][0]:xyxy[0][2]]
                print(count, crop_img.shape[0] * crop_img.shape[1])
                if crop_img.shape[0] < 414 and crop_img.shape[1] < 414:
                    # res = Watershed(crop_img)
                    # cv2.imshow('box', res)
                    cv2.imshow('box2', crop_img)
                    # cv2.rectangle(img, (xyxy[0][0], xyxy[0][1]), (xyxy[0][2], xyxy[0][3]), (0, 0, 255), 5)
                    # cv2.imshow('box', image_resize(img, 0.2))
                    # cv2.waitKey()
                    print(count)
                    cv2.imwrite(
                        '/home/yimu/Backup/projects_zy/datasets/hangtian_kegong/dianrong_noclass/images/test_save/' + str(
                            count) + '.jpg', crop_img)
                    count += 1
        # print(result[0])
